#!/usr/bin/env python3
"""
使用Kimi AI对政策文章进行分类的工具模块
"""

import json
import logging
import requests
from config.settings import KIMI_API_KEY, KIMI_BASE_URL, KIMI_MODEL, CATEGORIES

# 配置日志
logger = logging.getLogger(__name__)


class KimiAIClassifier:
    """Kimi AI文章分类器"""
    
    def __init__(self):
        """初始化Kimi AI配置"""
        self.api_key = KIMI_API_KEY
        self.base_url = KIMI_BASE_URL
        self.model = KIMI_MODEL
        self.categories = CATEGORIES
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
    def classify_article(self, title, content):
        """
        使用Kimi AI对文章进行分类
        
        Args:
            title (str): 文章标题
            content (str): 文章正文
            
        Returns:
            list: 匹配的分类列表
        """
        try:
            # 构造提示词
            prompt = f"""
            请根据以下文章的标题和正文内容，判断它属于哪个或哪些类别。
            
            可选的类别有：
            {', '.join(self.categories)}
            
            文章标题：{title}
            
            文章正文：
            {content[:30000]}  # 使用更长的内容长度以充分利用128k模型的能力
            
            请仅返回匹配的类别名称，多个类别用逗号分隔。如果没有匹配的类别，请返回"其他"。
            """
            
            # 构造请求数据
            data = {
                "model": self.model,
                "messages": [
                    {
                        "role": "system",
                        "content": "你是Kimi AI，一个专业的政策分析助手。你的任务是根据文章内容准确分类。"
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                "temperature": 0.1
            }
            
            # 调用Kimi AI API
            response = requests.post(
                f"{self.base_url}/chat/completions",
                headers=self.headers,
                json=data,
                timeout=30
            )
            
            response.raise_for_status()
            result = response.json()
            
            # 解析响应
            response_text = result['choices'][0]['message']['content'].strip()
            logger.info(f"Kimi AI分类结果: {response_text}")
            
            # 处理响应结果
            if ',' in response_text:
                # 多个分类
                categories = [cat.strip() for cat in response_text.split(',')]
            else:
                # 单个分类
                categories = [response_text.strip()]
                
            # 验证分类是否在预定义列表中
            valid_categories = [cat for cat in categories if cat in self.categories]
            
            # 如果没有有效分类，则归为"其他"
            if not valid_categories:
                valid_categories = ['其他']
                
            return valid_categories
            
        except Exception as e:
            logger.error(f"使用Kimi AI分类文章时出错: {str(e)}")
            # 出错时默认归为"其他"类
            return ['其他']
            
    def classify_articles_batch(self, articles):
        """
        批量分类文章
        
        Args:
            articles (list): 文章列表，每个文章应包含'title'和'content'字段
            
        Returns:
            list: 带有分类信息的文章列表
        """
        classified_articles = []
        
        for article in articles:
            title = article.get('标题', '') or article.get('title', '')
            content = article.get('正文', '') or article.get('content', '')
            
            # 获取分类
            categories = self.classify_article(title, content)
            
            # 添加分类信息到文章
            article['分类'] = categories
            classified_articles.append(article)
            
        return classified_articles